 package com.example.scribe.backend;

import com.example.scribe.SyntaxNode;
import net.bytebuddy.implementation.bytecode.StackManipulation;
import com.example.scribe.backend.handlers.*;

/**
 * Basic implementation of ExpressionEmitter using Strategy pattern with handlers.
 * Each expression type is handled by a dedicated handler registered in the registry.
 */
public final class BasicExpressionEmitter implements ExpressionEmitter {
    private final NodeHandlerRegistry registry;
    
    BasicExpressionEmitter() {
        this.registry = new NodeHandlerRegistry();
        initializeHandlers();
    }
    
    /**
     * Initialize and register all expression handlers.
     */
    private void initializeHandlers() {
        // Register handlers that don't need exprEmitter reference
        registry.registerExpressionHandler("identifier", new IdentifierExpressionHandler());
        registry.registerExpressionHandler("Identifier", new IdentifierExpressionHandler());
        registry.registerExpressionHandler("Terminal:NUMBER", new NumberExpressionHandler());
        
        // Register handlers that need exprEmitter (will be set with self-reference)
        registry.registerExpressionHandler("Call", new CallExpressionHandler(this));
        registry.registerExpressionHandler("addition", new ArithmeticExpressionHandler(this));
        registry.registerExpressionHandler("subtraction", new ArithmeticExpressionHandler(this));
        registry.registerExpressionHandler("multiplication", new ArithmeticExpressionHandler(this));
        registry.registerExpressionHandler("bitwise_and", new BitwiseExpressionHandler(this));
        registry.registerExpressionHandler("bitwise_or", new BitwiseExpressionHandler(this));
        registry.registerExpressionHandler("bitwise_xor", new BitwiseExpressionHandler(this));
        registry.registerExpressionHandler("shift", new BitwiseExpressionHandler(this));
        registry.registerExpressionHandler("unary", new UnaryExpressionHandler(this));
        registry.registerExpressionHandler("RangeExpr", new RangeExpressionHandler(this));
        registry.registerExpressionHandler("equality", new ComparisonExpressionHandler(this));
        registry.registerExpressionHandler("comparison", new ComparisonExpressionHandler(this));
    }

    @Override
    public StackManipulation emit(SyntaxNode expr, String generatedClassName, LocalVariableAllocator locals) {
        if (expr == null) return StackManipulation.Trivial.INSTANCE;
        
        // Top-level check: direct identifier or number nodes
        String exprType = expr.type();
        
        // Try handler for direct node type
        ExpressionNodeHandler handler = registry.getExpressionHandler(exprType);
        if (handler != null && handler.canHandle(exprType)) {
            StackManipulation result = handler.handle(expr, generatedClassName, locals);
            if (result != StackManipulation.Trivial.INSTANCE) {
                return result;
            }
        }
        
        // Handle identifier with locals check (even if handler returned Trivial)
        if (ASTUtils.isIdentifier(expr) && expr.text() != null && locals.hasName(expr.text())) {
            return BytecodeOps.loadInt(locals.getIndexByName(expr.text()));
        }
        
        // Handle number literal
        if (ASTUtils.isNumber(expr)) {
            try {
                int v = Integer.parseInt(expr.text());
                return BytecodeOps.pushInt(v);
            } catch (NumberFormatException e) {
                // continue
            }
        }
        
        // Prioritize structured operations (shift, unary, etc.) over simple operations
        // First, recursively search for structured operation nodes (they may be deeply nested)
        StackManipulation structured = findStructuredOperationRecursive(expr, generatedClassName, locals);
        if (structured != StackManipulation.Trivial.INSTANCE) {
            return structured;
        }
        
        // Then check direct children for structured operations
        for (SyntaxNode child : expr.children()) {
            String childType = child.type();
            // First check for structured operation nodes (shift, unary have higher priority)
            if ("shift".equals(childType) || "unary".equals(childType) ||
                "bitwise_or".equals(childType) || "bitwise_xor".equals(childType) ||
                "bitwise_and".equals(childType) || "addition".equals(childType) ||
                "multiplication".equals(childType) || "comparison".equals(childType) ||
                "equality".equals(childType)) {
                StackManipulation nested = emit(child, generatedClassName, locals);
                if (nested != StackManipulation.Trivial.INSTANCE) return nested;
            }
        }
        
        // Process children - look for Call, primary, arithmetic, comparison nodes
        for (SyntaxNode child : expr.children()) {
            String childType = child.type();
            
            // Try handler for child node type
            ExpressionNodeHandler childHandler = registry.getExpressionHandler(childType);
            if (childHandler != null && childHandler.canHandle(childType)) {
                StackManipulation result = childHandler.handle(child, generatedClassName, locals);
                if (result != StackManipulation.Trivial.INSTANCE) {
                    return result;
                }
            }
            
            // Handle primary nodes with identifier or number
            if ("primary".equals(childType)) {
                for (SyntaxNode prim : child.children()) {
                    if (ASTUtils.isNumber(prim)) {
                        int v = Integer.parseInt(prim.text());
                        return BytecodeOps.pushInt(v);
                    }
                    if (ASTUtils.isIdentifier(prim)) {
                        String name = prim.text();
                        if (name != null && locals.hasName(name)) {
                            int idx = locals.getIndexByName(name);
                            return BytecodeOps.loadInt(idx);
                        }
                        // not a local/param; keep scanning to allow outer Call handling
                    }
                }
            }
        }
        
        // Then check for direct identifier/number nodes that might be wrapped
        for (SyntaxNode child : expr.children()) {
            // Direct check for identifier as local variable
            if (ASTUtils.isIdentifier(child) && child.text() != null && locals.hasName(child.text())) {
                return BytecodeOps.loadInt(locals.getIndexByName(child.text()));
            }
            // Direct check for number
            if (ASTUtils.isNumber(child)) {
                try {
                    int v = Integer.parseInt(child.text());
                    return BytecodeOps.pushInt(v);
                } catch (NumberFormatException e) {
                    // continue
                }
            }
            // Try recursive emit
            StackManipulation nested = emit(child, generatedClassName, locals);
            if (nested != StackManipulation.Trivial.INSTANCE) return nested;
            
            // If emit returned Trivial, try deep recursive search for identifier/number in child's subtree
            StackManipulation deep = findValueEmit(child, generatedClassName, locals);
            if (deep != StackManipulation.Trivial.INSTANCE) return deep;
        }
        
        return StackManipulation.Trivial.INSTANCE;
    }

    /**
     * Recursively search for structured operation nodes (shift, unary, etc.) in the expression tree.
     * This method directly calls emit on structured operation nodes to avoid infinite recursion.
     */
    private StackManipulation findStructuredOperationRecursive(SyntaxNode node, String generatedClassName, LocalVariableAllocator locals) {
        if (node == null) {
            return StackManipulation.Trivial.INSTANCE;
        }

        String nodeType = node.type();
        // Check if this node is a structured operation - directly emit it using handler
        // Priority: comparison/equality have highest priority (used in if conditions)
        // Then bitwise operations, then shift and unary
        if ("comparison".equals(nodeType) || "equality".equals(nodeType)) {
            // Only handle comparison/equality nodes with operator (3+ children: left, operator, right)
            if (node.children().size() >= 3) {
                // Directly call handler without going through emit to avoid recursion
                ExpressionNodeHandler handler = registry.getExpressionHandler(nodeType);
                if (handler != null && handler.canHandle(nodeType)) {
                    StackManipulation result = handler.handle(node, generatedClassName, locals);
                    if (result != StackManipulation.Trivial.INSTANCE) {
                        return result;
                    }
                }
            }
            // If comparison/equality node has no operator, continue searching children
            // Don't return Trivial here, let it fall through to recursive search
        } else if ("bitwise_or".equals(nodeType) || "bitwise_and".equals(nodeType) || "bitwise_xor".equals(nodeType)) {
            // Only handle bitwise nodes with operator (3+ children: left, operator, right)
            if (node.children().size() >= 3) {
                // Directly call handler without going through emit to avoid recursion
                ExpressionNodeHandler handler = registry.getExpressionHandler(nodeType);
                if (handler != null && handler.canHandle(nodeType)) {
                    StackManipulation result = handler.handle(node, generatedClassName, locals);
                    if (result != StackManipulation.Trivial.INSTANCE) {
                        return result;
                    }
                }
            }
            // If bitwise node has no operator, continue searching children
            // Don't return Trivial here, let it fall through to recursive search
        } else if ("shift".equals(nodeType)) {
            // Only handle shift nodes with operator (3+ children: left, operator, right)
            if (node.children().size() >= 3) {
                // Directly call handler without going through emit to avoid recursion
                ExpressionNodeHandler handler = registry.getExpressionHandler(nodeType);
                if (handler != null && handler.canHandle(nodeType)) {
                    StackManipulation result = handler.handle(node, generatedClassName, locals);
                    if (result != StackManipulation.Trivial.INSTANCE) {
                        return result;
                    }
                }
            }
            // If shift node has no operator (only 1 child), continue searching children
            // Don't return Trivial here, let it fall through to recursive search
        } else if ("unary".equals(nodeType)) {
            // Only handle unary nodes with operator (2+ children: operator, operand)
            // Check if first child is an operator terminal
            boolean hasOperator = false;
            if (node.children().size() >= 2) {
                SyntaxNode firstChild = node.children().get(0);
                if (firstChild.type().startsWith("Terminal:")) {
                    String text = firstChild.text();
                    if ("!".equals(text) || "-".equals(text) || "~".equals(text)) {
                        hasOperator = true;
                    }
                }
            }
            
            if (hasOperator) {
                // Directly call handler without going through emit to avoid recursion
                ExpressionNodeHandler handler = registry.getExpressionHandler(nodeType);
                if (handler != null && handler.canHandle(nodeType)) {
                    StackManipulation result = handler.handle(node, generatedClassName, locals);
                    if (result != StackManipulation.Trivial.INSTANCE) {
                        return result;
                    }
                }
            }
            // If unary node has no operator, continue searching children
            // Don't return Trivial here, let it fall through to recursive search
        }

        // Recursively search children
        // Important: For intermediate nodes like assignment, conditional, logical_or, logical_and, equality,
        // we need to recursively search to find comparison nodes that may be deeply nested
        for (SyntaxNode child : node.children()) {
            // Skip operator terminals
            if (!child.type().startsWith("Terminal:")) {
                StackManipulation found = findStructuredOperationRecursive(child, generatedClassName, locals);
                if (found != StackManipulation.Trivial.INSTANCE) {
                    return found;
                }
            }
        }

        return StackManipulation.Trivial.INSTANCE;
    }

    /**
     * Helper method to find and emit value nodes (identifier or number).
     * Used for nested node searches when direct emit returns Trivial.
     */
    private StackManipulation findValueEmit(SyntaxNode node, String generatedClassName, LocalVariableAllocator locals) {
        if (node == null) return StackManipulation.Trivial.INSTANCE;
        
        // For structured nodes (comparison, equality, shift, unary, etc.), use emit directly
        if ("comparison".equals(node.type()) || "equality".equals(node.type()) ||
            "shift".equals(node.type()) || "unary".equals(node.type()) ||
            "bitwise_or".equals(node.type()) || "bitwise_xor".equals(node.type()) ||
            "bitwise_and".equals(node.type()) || "addition".equals(node.type()) ||
            "multiplication".equals(node.type())) {
            StackManipulation result = emit(node, generatedClassName, locals);
            if (result != StackManipulation.Trivial.INSTANCE) return result;
        }
        
        // For expression nodes, first check if it contains structured operations
        if ("expression".equals(node.type())) {
            // Look for structured operation nodes first
            for (SyntaxNode child : node.children()) {
                String childType = child.type();
                if ("comparison".equals(childType) || "equality".equals(childType) ||
                    "shift".equals(childType) || "unary".equals(childType) ||
                    "bitwise_or".equals(childType) || "bitwise_xor".equals(childType) ||
                    "bitwise_and".equals(childType) || "addition".equals(childType) ||
                    "multiplication".equals(childType)) {
                    StackManipulation result = emit(child, generatedClassName, locals);
                    if (result != StackManipulation.Trivial.INSTANCE) return result;
                }
            }
            // If no structured node found, continue searching for other nodes
        }
        
        // Direct check for leaf nodes
        if (ASTUtils.isNumber(node)) {
            try {
                int v = Integer.parseInt(node.text());
                return BytecodeOps.pushInt(v);
            } catch (NumberFormatException e) {
                return StackManipulation.Trivial.INSTANCE;
            }
        }
        if (ASTUtils.isIdentifier(node)) {
            if (locals.hasName(node.text())) {
                return BytecodeOps.loadInt(locals.getIndexByName(node.text()));
            }
            // Not a local/param - might be function name, continue searching children
        }
        
        // Try direct emit for other structured nodes (Call, addition, etc.)
        StackManipulation result = emit(node, generatedClassName, locals);
        if (result != StackManipulation.Trivial.INSTANCE) return result;
        
        // Recursively search children
        for (SyntaxNode child : node.children()) {
            StackManipulation nested = findValueEmit(child, generatedClassName, locals);
            if (nested != StackManipulation.Trivial.INSTANCE) return nested;
        }
        return StackManipulation.Trivial.INSTANCE;
    }
}
